import subprocess

from configparser import ConfigParser

def update_ini_file_from_command(ini_file: str, EPOCH: int, LR: float, NAME: str):
    # 创建ConfigParser对象
    config = ConfigParser()

    # 尝试读取已有的配置文件，如果文件不存在则会新建
    config.read(ini_file)
    #TODO
    # 模型部分（model）
    if not config.has_section("model"):
        config.add_section("model")
    #TODO
    config.set("model", "pretrained_model_name_or_path", "your_model_path")

    # 数据部分（data）
    if not config.has_section("data"):
        config.add_section("data")
    config.set("data", "train_dataset_path", f"../../data/test_set_B/{NAME}.parquet")
    config.set("data", "eval_dataset_path", f"../../data/test_set_B/{NAME}.parquet")
    config.set("data", "max_length", "2048")

    # 训练部分（training）
    if not config.has_section("training"):
        config.add_section("training")
    config.set("training", "run_name", f"{NAME}_e{EPOCH}_l{LR}")
    config.set("training", "num_train_epochs", str(EPOCH))
    config.set("training", "learning_rate", str(LR))
    config.set("training", "output_dir", f"../../checkpoints/{NAME}_e{EPOCH}_l{LR}")
    config.set("training", "eval_steps", "50")
    config.set("training", "save_steps", "100")

    config.set("training", "per_device_train_batch_size", "1")
    config.set("training", "per_device_eval_batch_size", "1")

    # 打印更新后的配置
    print("Updated config sections:", dict(config._sections))

    # 保存修改后的 ini 文件
    with open(ini_file, 'w') as configfile:
        config.write(configfile)

    print(f"Config file {ini_file} updated successfully from command.")


def run_accelerate_train():
    command = [
        "accelerate", "launch",
        "--config_file", "config.yaml",
        "--main_process_port", "29507",
        "train.py"
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    if result.returncode == 0:
        print("Training started successfully!")
        print(f"Output: {result.stdout}")  # 打印标准输出
    else:
        print(f"Error: {result.stderr}")  # 打印错误输出
        print(f"Command failed with return code {result.returncode}")
        print(f"Stdout: {result.stdout}")
        print(f"Stderr: {result.stderr}")

# 定义参数
EPOCH = 10
LR = 5e-4
NAME = "mix_cl_a50"
ini_file = "train_config.ini"

# 更新 ini 配置文件
update_ini_file_from_command(ini_file, EPOCH, LR, NAME)

from model_load import train_mod


train_mod()  # 调用 main.py 中的 main 函数

# run_accelerate_train()
